"""
AI代码生成工作流演示脚本

此脚本演示了一个完整的AI辅助代码生成workflow，包含以下步骤：

1. 算法分析阶段：
   - 接收用户输入的算法描述（如"Calculate Fibonacci sequence"）
   - 使用AI分析算法需求，提供函数签名和输入输出示例

2. 伪代码生成阶段：
   - 基于算法分析结果生成结构化的伪代码
   - 伪代码包含输入输出定义和算法逻辑描述

3. Python代码生成阶段：
   - 将伪代码转换为可执行的Python代码
   - 确保代码遵循伪代码的逻辑结构

4. 代码提取与执行阶段：
   - 从AI回复中提取代码块
   - 执行生成的代码进行验证
   - 将代码保存到文件中

整个workflow展示了从自然语言描述到可执行代码的完整转换过程。
"""

from utils.api_connector import ApiConnector
import re

def extract_code_block(text):
    """
    提取文本中的代码块，支持语言标识符
    返回第一个找到的代码块，如果没有找到则返回None
    """
    # 使用正则表达式匹配代码块，支持可选的语言标识符
    pattern = r'```(?:\w+)?\s*\n?(.*?)```'
    matches = re.findall(pattern, text, re.DOTALL)
    
    if matches:
        # 返回第一个匹配的代码块，去除首尾空白
        return matches[0].strip()
    return None

if __name__ == "__main__":
    analyzer = ApiConnector()
    
    userinput = "Calculate Fibonacci sequence"
    template = "Analyze the algorithm described in the user's input. " \
    "Provide the function signature including its inputs and outputs, and give " \
    "exactly 3 input-output examples to demonstrate its behavior. " \
    "Do not include any additional commentary. " \
    "User input: {userinput}"
    prompt = template.format(userinput=userinput)

    reply, reasoning = analyzer.send_request(userinput, verbose=False)

    print(f"Reply: {reply}")
    print("================================================")


    template = """Generate pseudocode for the algorithm described in the user's input. 
The pseudocode must begin with the definition of input and output. 
The code should be enclosed within triple backticks like ``` ```.
User input: {userinput}"""
    prompt = template.format(userinput=reply)
    print(f"Prompt: {prompt}")
    reply, reasoning = analyzer.send_request(prompt, verbose=False)
    
    print(f"Reply: {reply}")
    print("================================================")

    # 使用改进的代码块提取函数
    pseudocode = extract_code_block(reply)
    if pseudocode:
        print("Extracted Pseudocode:")
        print(pseudocode)
    else:
        print("No pseudocode found in the reply.")
    print("================================================")

    template = "Generate Python code for the algorithm described in the user's input. " \
    "Follow the pseudocode provided in the previous step. The code should be enclosed within triple backticks like ``` ```." \
    "User input: {userinput}"
    prompt = template.format(userinput=pseudocode)
    print(f"Prompt: {prompt}")
    reply, reasoning = analyzer.send_request(prompt, verbose=False)
    print(f"Reply: {reply}")
    print("================================================")

    # 使用相同的函数提取Python代码
    code = extract_code_block(reply)
    if code:
        print("Extracted Code:")
        print(code)
    else:
        print("No code found in the reply.")
    print("================================================")

    # Execute the extracted Python code
    try:
        exec(code)
    except Exception as e:
        print(f"Error while executing the code: {e}")
    
    # write the code to a file
    with open("fibonacci.py", "w") as f:
        f.write(code)
        print("Code written to fibonacci.py")
